import torch
import torchvision.transforms as transforms
from PIL import Image
import torch.nn as nn
from torchvision import models

# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 定义并加载模型
alexnet = models.alexnet()
alexnet.classifier[6] = nn.Linear(alexnet.classifier[6].in_features, 3)
alexnet.load_state_dict(torch.load('alexnet_best_model.pth', map_location=device))
alexnet = alexnet.to(device)
alexnet.eval()

# 图像预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])

# 预测函数
def predict(image_path):
    # 加载并预处理图像
    image = Image.open(image_path).convert('RGB')
    image = transform(image)
    image = image.unsqueeze(0).to(device)  # 增加批次维度并移至设备

    # 模型预测
    with torch.no_grad():
        outputs = alexnet(image)
        _, predicted = torch.max(outputs, 1)
        predicted = predicted.cpu().numpy()

    # 类别名称
    class_names = ['最优级-1级', '中间级-2级', '差等级-3级']  # 类别名
    return class_names[predicted[0]]

# # 示例用法
# image_path = r'D:\develop\PythonCode\python基础\附_项目实战\九_薄膜图片级别分类\test_img\test_605.tif'  # 替换为实际的图片路径
# predicted_class = predict(image_path)
# print("Predicted Class:", predicted_class)
image_path_list = [
    r'D:\develop\PythonCode\python基础\附_项目实战\九_薄膜图片级别分类\data\real_data\level_1\1.tif',
    r'D:\develop\PythonCode\python基础\附_项目实战\九_薄膜图片级别分类\data\real_data\level_2\1.tif',
    r'D:\develop\PythonCode\python基础\附_项目实战\九_薄膜图片级别分类\data\real_data\level_3\1.tif'

]  # 替换为你的图片路径

for image_path in image_path_list:
    predicted_class = predict(image_path)
    print("Predicted Class:", predicted_class)